import pandas as pd
from imblearn.over_sampling import RandomOverSampler, SMOTE
from imblearn.under_sampling import RandomUnderSampler
from sklearn.compose import ColumnTransformer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, roc_auc_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder, StandardScaler


def train_and_evaluate(model, X_train, y_train, X_test, y_test):
    """
    训练并评估模型
    """
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    auc = roc_auc_score(y_test, y_pred)
    report = classification_report(y_test, y_pred)
    return auc, report


# 读取训练集数据
df = pd.read_csv(r'D:\WorkArea\WorkSpace\Python\talents_loss\data\raw\train.csv')

print('数据基本信息：')
df.info()

# 查看数据集行数和列数
rows, columns = df.shape

if rows < 100 and columns < 20:
    # 短表数据（行数少于100且列数少于20）查看全量数据信息
    print('数据全部内容信息：')
    print(df.to_csv(sep='\t', na_rep='nan'))
else:
    # 长表数据查看数据前几行信息
    print('数据前几行内容信息：')
    print(df.head().to_csv(sep='\t', na_rep='nan'))

# 划分特征和目标变量
X = df.drop('Attrition', axis=1)
y = df['Attrition']

# 对 object 类型数据进行独热编码，对数值型数据进行标准化
categorical_cols = X.select_dtypes(include=['object']).columns
numerical_cols = X.select_dtypes(exclude=['object']).columns

preprocessor = ColumnTransformer(
    transformers=[
        ('num', StandardScaler(), numerical_cols),
        ('cat', OneHotEncoder(), categorical_cols)
    ])

X = preprocessor.fit_transform(X)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42,stratify=y)

# 原始数据训练逻辑回归模型
logreg = LogisticRegression(random_state=42)
auc, report = train_and_evaluate(logreg, X_train, y_train, X_test, y_test)
print('原始数据模型评估：')
print(f"AUC: {auc:.4f}")
print(report)

# 随机过采样
ros = RandomOverSampler(random_state=42)
X_resampled, y_resampled = ros.fit_resample(X_train, y_train)

logreg = LogisticRegression(random_state=42)
auc, report = train_and_evaluate(logreg, X_resampled, y_resampled, X_test, y_test)
print('随机过采样模型评估：')
print(f"AUC: {auc:.4f}")
print(report)

# SMOTE 过采样
smote = SMOTE(random_state=42)
X_resampled, y_resampled = smote.fit_resample(X_train, y_train)

logreg = LogisticRegression(random_state=42)
auc, report = train_and_evaluate(logreg, X_resampled, y_resampled, X_test, y_test)
print('SMOTE 过采样模型评估：')
print(f"AUC: {auc:.4f}")
print(report)

# 随机欠采样
rus = RandomUnderSampler(random_state=42)
X_resampled, y_resampled = rus.fit_resample(X_train, y_train)

logreg = LogisticRegression(random_state=42)
auc, report = train_and_evaluate(logreg, X_resampled, y_resampled, X_test, y_test)
print('随机欠采样模型评估：')
print(f"AUC: {auc:.4f}")
print(report)

# 调整类别权重
logreg_weighted = LogisticRegression(random_state=42, class_weight='balanced')
auc, report = train_and_evaluate(logreg_weighted, X_train, y_train, X_test, y_test)
print('调整类别权重模型评估：')
print(f"AUC: {auc:.4f}")
print(report)